from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import KFold, train_test_split
from sklearn.metrics import roc_auc_score
from catboost import CatBoostClassifier
from base import base
class LR(base):
    def __init__(self, cfg, mode='train') -> None:
        super().__init__(cfg)
        # self.clf = LogisticRegression()
        self.mode = mode
        self.seed = self.train_cfg.get("seed")
        params = {'learning_rate': 0.15, 'depth': 6, 'l2_leaf_reg': 20, 'bootstrap_type':'Bernoulli','random_seed':self.seed,
                'od_type': 'Iter', 'od_wait': 50, 'random_seed': 11, 'allow_writing_files': False}
        # self.kf = KFold(n_splits=5, shuffle=True, random_state=self.seed)

        self.clf = CatBoostClassifier(iterations=1000, **params, eval_metric='AUC')

    def __call__(self, x):
        train_x = x.drop(['index', 'label'], axis=1)
        train_y = x['label']
        print(x.columns)
        if self.mode == "train":
            test_size = self.train_cfg.get('validation_split')

            train_x, val_x, train_y, val_y = train_test_split(train_x, train_y, test_size=test_size, random_state=self.seed)
            self.clf.fit(train_x, train_y, eval_set=(val_x, val_y),
                    metric_period=200,
                    cat_features=[], 
                    use_best_model=True, 
                    verbose=1)
            val_pred  = self.clf.predict_proba(val_x)[:,1]
            auc_score = roc_auc_score(val_y, val_pred)
        else:
            self.clf.fit(train_x, train_y,
                    metric_period=200,
                    cat_features=[], 
                    use_best_model=False, 
                    verbose=1)
            train_pred  = self.clf.predict_proba(train_x)[:,1]
            auc_score = roc_auc_score(train_y, train_pred)
            
        print("==================auc:{}=========================".format(auc_score))


        
    def test(self, x):
        test_x = x.drop(['index'], axis=1)
        x['predict'] = self.clf.predict_proba(test_x)[:, 1]
        return x[['index', 'predict']]
    

